import os

import numpy as np
import cv2

img_path = r'./data/train/imgs'
img_names = os.listdir(img_path)
img_names.sort()

mask_path = r'./data/train/masks'
mask_names = os.listdir(mask_path)
mask_names.sort()


# # 翻转
# for i, (file,mask) in enumerate(zip(img_names, mask_names)):
#     img = cv2.imread(os.path.join(img_path, file))
#     mask = cv2.imread(os.path.join(mask_path, mask))
#     img = cv2.flip(img, -1)
#     mask = cv2.flip(mask, -1)
#     cv2.imwrite(r'./data/train/imgs/inc{}.jpg'.format(i), img)
#     cv2.imwrite(r'./data/train/masks/inc{}.bmp'.format(i), mask)
# print('done')

# # 随机旋转
# for i, (file,mask) in enumerate(zip(img_names, mask_names)):
#     angle = np.random.randint(-30,30)
#     img = cv2.imread(os.path.join(img_path, img_names[i]))
#     mask = cv2.imread(os.path.join(mask_path, mask_names[i]), 0)  # mask灰度图
#     H,W = img.shape[:2]
#     M = cv2.getRotationMatrix2D(center=(W//2,H//2), angle=angle, scale=1.0)
#     img = cv2.img = warpAffine(img, M, dsize=(W,H), dst=None)
#     mask = cv2.warpAffine(mask, M, dsize=(W,H),dst=None)
#     cv2.imwrite(r'./data/train/imgs/rota{}.jpg'.format(i), img)
#     cv2.imwrite(r'./data/train/masks/rota{}.bmp'.format(i), mask)
# print('done')

# 随机缩放
for i, (file, mask) in enumerate(zip(img_names, mask_names)):
    scale = np.random.random() + 0.5
    img = cv2.imread(os.path.join(img_path, img_names[i]))
    mask = cv2.imread(os.path.join(mask_path, mask_names[i]), 0)
    img = cv2.resize(img, None, None, scale, scale)
    mask = cv2.resize(mask, None, None, scale, scale)
    cv2.imwrite(r'./data/train/imgs/resize{}.jpg'.format(i), img)
    cv2.imwrite(r'./data/train/masks/resize{}.bmp'.format(i), mask)
print('done')